import torch
torch.set_num_threads(1)
import torch.nn as nn
import torch.nn.functional as F
from src.models.layer.embedding_layer import TabularEmbedding

class DepthwiseSeparableConv1d(nn.Module):
    """1D深度可分离卷积，适配表格数据"""
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super(DepthwiseSeparableConv1d, self).__init__()
        self.depthwise = nn.Conv1d(in_channels, in_channels, kernel_size, stride, padding, groups=in_channels)
        self.pointwise = nn.Conv1d(in_channels, out_channels, 1)
        self.bn = nn.BatchNorm1d(out_channels)
        self.relu = nn.ReLU6(inplace=True)
        
    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

class MobileNetBlock(nn.Module):
    """MobileNet块，用于表格数据特征提取"""
    def __init__(self, in_channels, out_channels, stride=1, expand_ratio=6):
        super(MobileNetBlock, self).__init__()
        self.stride = stride
        self.use_residual = stride == 1 and in_channels == out_channels
        
        # 扩展层
        expanded_channels = in_channels * expand_ratio
        self.expand = nn.Sequential(
            nn.Conv1d(in_channels, expanded_channels, 1, bias=False),
            nn.BatchNorm1d(expanded_channels),
            nn.ReLU6(inplace=True)
        ) if expand_ratio != 1 else None
        
        # 深度可分离卷积
        self.depthwise = nn.Conv1d(
            expanded_channels if expand_ratio != 1 else in_channels,
            expanded_channels if expand_ratio != 1 else in_channels,
            3, stride, 1, 
            groups=expanded_channels if expand_ratio != 1 else in_channels,
            bias=False
        )
        self.bn1 = nn.BatchNorm1d(expanded_channels if expand_ratio != 1 else in_channels)
        self.relu1 = nn.ReLU6(inplace=True)
        
        # 压缩层
        self.pointwise = nn.Conv1d(
            expanded_channels if expand_ratio != 1 else in_channels,
            out_channels, 1, bias=False
        )
        self.bn2 = nn.BatchNorm1d(out_channels)
        
    def forward(self, x):
        residual = x
        
        if self.expand is not None:
            x = self.expand(x)
        
        x = self.depthwise(x)
        x = self.bn1(x)
        x = self.relu1(x)
        
        x = self.pointwise(x)
        x = self.bn2(x)
        
        if self.use_residual:
            x = x + residual
            
        return x

class MobileNetClassifier(nn.Module):
    """适配表格数据的MobileNet模型，使用TabularEmbedding"""
    def __init__(self, num_cat_features, num_cont_features, num_classes=2, 
                 embedding_dim=16, cont_hidden_dim=128, hidden_dim=128, dropout=0.5):
        super(MobileNetClassifier, self).__init__()
        
        # 表格嵌入层
        self.embedding = TabularEmbedding(
            num_cat_features=num_cat_features,
            num_cont_features=num_cont_features,
            embedding_dim=embedding_dim,
            cont_hidden_dim=cont_hidden_dim,
            fuse=True  # 直接融合类别和连续特征
        )
        
        # 融合后的总特征维度
        total_features = num_cat_features * embedding_dim + cont_hidden_dim
        
        # 将1D特征转换为2D以便使用卷积
        self.feature_reshape = nn.Linear(total_features, hidden_dim * 4)
        
        # MobileNet特征提取（简化版）
        self.features = nn.Sequential(
            nn.Conv1d(1, 32, 3, 1, 1, bias=False),
            nn.BatchNorm1d(32),
            nn.ReLU6(inplace=True),
            MobileNetBlock(32, 64, stride=2),
            MobileNetBlock(64, 64),
            nn.AdaptiveAvgPool1d(1)
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(64, hidden_dim),
            nn.ReLU6(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes)
        )
        
    def forward(self, batch):
        x_cat, x_cont, _ = batch
        
        # 使用TabularEmbedding融合类别和连续特征
        x_combined, _, _ = self.embedding(x_cat, x_cont)
        
        # 重塑为卷积输入格式
        x_reshaped = self.feature_reshape(x_combined)
        x_reshaped = x_reshaped.unsqueeze(1)  # [batch_size, 1, features]
        
        # MobileNet特征提取
        x_features = self.features(x_reshaped)
        x_features = x_features.squeeze(-1)  # [batch_size, channels]
        
        # 分类
        output = self.classifier(x_features)
        return output